#%%
import pandas

#%%
def SercherCode(df, inputstr, module=0):
    """
    + inputstr: 关键词
    + module = 列名
        + {0: 'ts_code' 代码含区域
        + 1: 'name' 股票中文名字
        + 2: 'area' 上市区域
        + 3: 'industry' 行业类目
        + 4: 'market' 上市板块(主板 三板 创业板)
        + 5: 'exchange' 区域
        + 6: 'list_date' 上市时间 inputstr = starttime, endtime
    return
    pandas.Dataframe
            ts_code name area industry market exchange  list_date
    1218  002785.SZ  万里石   福建     其他建材    中小板     SZSE   20151223
    """
    # 创建列表序列
    if isinstance(inputstr, int) or module == 0:
        inputstr = str(inputstr)
    elif isinstance(inputstr, list) or module == 6:
        
        info = SerchbyTime(df, starTime=0, endTime=20201221)
        return info
    namelist = df
    dic = {}
    a = 0
    columns = namelist.columns.tolist()
    for i in columns:
        dic[a] = i
        a += 1
    # print("可选编号",dic)
    incolumns = dic[module]  # 需要手动输入
    info = namelist.loc[namelist[incolumns].str.contains(inputstr)]  # 模糊搜索
    print("\n" * 2)
    # print(info)
    return info


def SerchbyTime(df, starTime=0, endTime=20201221):
    """
    + SerchbyTime(20200101, 20201111)
    + 进行时间检索，找到某个时间段内的股票
    + 输入值为 int 20200101
    """
    namelist = df
    if starTime + endTime != 0:
        if starTime == 0:
            info = namelist[namelist["list_date"] < endTime]
        elif endTime == 0:
            info = namelist[namelist["list_date"] > starTime]
        elif starTime < endTime:
            info = namelist[
                (namelist["list_date"] > starTime) & (namelist["list_date"] < endTime)
            ]
        else:
            print("启始时间必须小于终止时间")
        return info
    else:
        print("启始时间与终止时间不可都为零")